import { StateGraph } from "@langchain/langgraph";
import { HumanMessage, AIMessage } from "@langchain/core/messages";
import axios from "axios";
import { Ollama } from "@langchain/ollama";

// --- 天气 API 工具 ---
async function getWeather(city) {
  const result = await axios.get(
    `${process.env.WEATHER_API_CITY_NAME}?location=${encodeURIComponent(
      city
    )}&key=${process.env.WEATHER_API_KEY}`
  );
  const location = result.data.location[0].id;
  const url = `${process.env.WEATHER_API_HOSTNAME}?location=${location}&key=${process.env.WEATHER_API_KEY}`;

  try {
    const response = await axios.get(url);
    return response.data;
  } catch (error) {
    console.error("天气API错误:", error.response?.data || error.message);
    return { error: "天气服务暂不可用" };
  }
}

// --- 初始化模型 ---
const model = new Ollama({
  baseUrl: process.env.MODEL_URL,
  model: "deepseek-r1:7b", // 建议使用中文模型
});

// --- 节点函数 ---
// 1. 提取城市名称
async function extractCity(state) {
  const { userInput, messages } = state;

  // 尝试从历史中获取上次查询的城市
  const lastCity = messages
    .slice()
    .reverse()
    .find((msg) => msg instanceof AIMessage && "city" in msg.additional_kwargs)
    ?.additional_kwargs?.city;

  // 简单正则匹配（实际可用NLP模型增强）
  const cityMatch =
    userInput.match(/^(.+?)(天气|下雨|晴天|温度)/)?.[1] ||
    userInput.match(/(北京|上海|广州|深圳)/)?.[0] ||
    lastCity ||
    "北京"; // 默认城市

  return {
    ...state,
    city: cityMatch,
    // 在消息中记录城市（供后续节点使用）
    messages: [
      ...messages,
      new HumanMessage(userInput),
      new AIMessage("", { city: cityMatch }), // 元数据存储城市
    ],
  };
}

// 2. 调用天气API
async function fetchWeather(state) {
  if (!state.city) throw new Error("未解析到城市名称");

  const weatherData = await getWeather(state.city);
  console.log(weatherData);
  return { ...state, weatherData };
}

// 3. 生成自然语言回复（流式）
async function generateResponse(state) {
  const { city, weatherData, messages } = state;

  if (weatherData.error) {
    return {
      ...state,
      messages: [...messages, new AIMessage(`抱歉，无法获取${city}的天气信息`)],
    };
  }

  const { temp, text } = weatherData.now;
  const prompt = `用户问题：${state.userInput}\n天气数据：${text}，温度 ${temp}°C`;

  // 流式生成回复
  const stream = await model.stream(prompt);

  let fullResponse = "";
  for await (const chunk of stream) {
    fullResponse += chunk;
    console.log(fullResponse);
  }

  return {
    ...state,
    streamingResponse: fullResponse,
    messages: [...messages, new AIMessage(fullResponse, { city, weatherData })],
  };
}

// --- 构建工作流 ---
const workflow = new StateGraph({
  channels: {
    messages: { value: (x) => x, default: () => [] },
    userInput: { value: (x) => x },
    city: { value: (x) => x },
    weatherData: { value: (x) => x },
    streamingResponse: { value: (x) => x },
  },
});
// 添加节点
const weatherAssistant = workflow
  .addNode("extract_city", extractCity)
  .addNode("fetch_weather", fetchWeather)
  .addNode("generate_response", generateResponse)
  .addEdge("__start__", "extract_city")
  .addEdge("extract_city", "fetch_weather")
  .addEdge("fetch_weather", "generate_response")
  .addEdge("generate_response", "__end__")
  .compile();

// --- 使用示例 ---
export async function sendMessage(userInput, history = []) {
  const initialState = {
    messages: history,
    userInput,
  };

  const result = await weatherAssistant.invoke(initialState);
  return {
    response: result.streamingResponse,
    messages: result.messages, // 更新后的对话历史
  };
}
